#!/usr/bin/python3
# coding=utf-8
#
# Copyright (C) 2025.
# Huawei Technologies Co., Ltd. All rights reserved.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# ===============================================================================

import numpy as np

dtype_emu = {np.float16: 1, np.float32: 0}


def gen_lp_norm_v2_data():
    # 配置输入
    dtype = np.float32
    input_shape = [11, 123]  # 可以随意修改测试
    axis = 1                 # 0=全局, 1=列规约, 2=行规约
    p_value = 2.0            # L2 范数

    rows, cols = input_shape
    total_length = rows * cols

    # 随机输入数据
    input_x = np.random.uniform(-10, 10, input_shape).astype(dtype)

    # Golden 计算
    if axis == 0:
        # 全局范数
        norm = np.linalg.norm(input_x.reshape(-1), ord=p_value)
        golden = input_x / norm
    elif axis == 1:
        # 按列规约 (axis=1 → 每列归一化)
        norm = np.linalg.norm(input_x, ord=p_value, axis=0, keepdims=True)
        golden = input_x / norm
    elif axis == 2:
        # 按行规约 (axis=2 → 每行归一化)
        norm = np.linalg.norm(input_x, ord=p_value, axis=1, keepdims=True)
        golden = input_x / norm
    else:
        raise ValueError("Unsupported axis = {}".format(axis))

    # 正确定义tiling结构体（与C++中SyncBNTrainingUpdateCustomTilingData对应）
    tiling = np.zeros((1), dtype=[
        ('totalLength', np.uint32),
        ('dataType',np.uint32),
        ('blockLength', np.uint32),
        ('tileNum', np.uint32),
        ('tileLength', np.uint32),
        ('lasttileLength', np.uint32),
        ('formerNum', np.uint32),
        ('formerLength', np.uint32),
        ('formerTileNum', np.uint32),
        ('formerTileLength', np.uint32),
        ('formerLastTileLength', np.uint32),
        ('tailNum', np.uint32),
        ('tailLength', np.uint32),
        ('tailTileNum', np.uint32),
        ('tailTileLength', np.uint32),
        ('tailLastTileLength', np.uint32),
        ('isEvenCore', np.uint32),

        ('pValue', np.float32),
        ('axis', np.uint32),
        ('cols', np.uint32),
        ('rows', np.uint32),
        ('epsilon', np.float32),
    ])

    tiling['totalLength'] = total_length  # 关键：设置总长度为11*123
    tiling['rows'] = rows
    tiling['cols'] = cols
    tiling['dataType'] = dtype_emu[dtype]
    tiling['epsilon'] = 1e-12
    tiling['axis'] = axis
    tiling['pValue']= p_value



    # 保存文件
    tiling.tofile("./input/input_tiling.bin")
    input_x.tofile("./input/input_x.bin")
    golden.astype(dtype).tofile("./output/golden.bin")

    print("✅ Data generated:")
    print("expect pnorm :",norm)
    print(" - input_x.bin shape:", input_x.shape, input_x.dtype)
    print(" - golden.bin shape:", golden.shape, golden.dtype)
    print(" - tiling:", tiling)


if __name__ == "__main__":
    gen_lp_norm_v2_data()
